#!/usr/bin/python3

import glob
import os
import string
import subprocess
import sys
import unittest
import errno
import time
import socket
import tempfile
import binascii
import logging

import importlib.machinery
import importlib.util

import parent
parent.ensure_bots()  # NOQA: testvm lives in bots/
import testlib
import testvm

sys.dont_write_bytecode = True
os.environ['PYTHONUNBUFFERED'] = '1'


def flush_stdout():
    while True:
        try:
            sys.stdout.flush()
            break
        except BlockingIOError:
            time.sleep(0.1)


class Test:
    def __init__(self, test_id, command, timeout, nondestructive, retry_when_affected):
        self.process = None
        self.retries = 0
        self.test_id = test_id
        self.command = command
        self.timeout = timeout
        self.nondestructive = nondestructive
        self.machine_id = None
        self.retry_when_affected = retry_when_affected
        self.returncode = None

    def assign_machine(self, machine_id, ssh_address, web_address):
        assert self.nondestructive, "assigning a machine only works for nondestructive test"
        self.machine_id = machine_id
        self.command.insert(-2, "--machine")
        self.command.insert(-2, ssh_address)
        self.command.insert(-2, "--browser")
        self.command.insert(-2, web_address)

    def start(self):
        if self.nondestructive:
            assert self.machine_id is not None, f"need to assign nondestructive test {self} {self.command} to a machine"
        self.outfile = tempfile.TemporaryFile()
        self.process = subprocess.Popen(["timeout", "-v", str(self.timeout)] + self.command,
                                        stdout=self.outfile, stderr=subprocess.STDOUT)

    def poll(self):
        poll_result = self.process.poll()
        if poll_result is not None:
            self.outfile.flush()
            self.outfile.seek(0)
            self.output = self.outfile.read()
            self.outfile.close()
            self.outfile = None
            self.returncode = self.process.returncode

        return poll_result

    def finish(self, affected_tests, print_tap=True, thorough=False, track_naughties=False):
        """Returns if a test should retry or not

        Call test-failure-policy on the test's output, print if needed.

        Return (retry_reason, exit_code). retry_reason can be None or a string.
        """

        affected = any([self.command[0].endswith(t) for t in affected_tests])
        retry_reason = ""
        skip_reason = ""

        # Try affected tests 3 times
        if self.returncode == 0 and affected and self.retry_when_affected and self.retries < 2:
            retry_reason = "test affected tests 3 times"
            self.retries += 1
            self._print_test(print_tap, "# RETRY {0} ({1})".format(self.retries, retry_reason))
            return retry_reason, 0

        # If test is being skipped pick up the reason
        if self.returncode == 77:
            lines = self.output.splitlines()
            skip_reason = lines[-1].strip().decode("utf-8")
            self.output = b"\n".join(lines[:-1])

        if self.returncode in [0, 77]:
            self._print_test(print_tap, skip_reason=skip_reason)
            return None, 0

        if not thorough:
            cmd = ["test-failure-policy"]
            if not track_naughties:
                cmd.append("--offline")
            cmd.append(testvm.DEFAULT_IMAGE)
            try:
                proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
                reason = proc.communicate(self.output + ("not ok " + str(self)).encode())[0].strip()

                if proc.returncode == 77:
                    self._print_test(skip_reason="# SKIP {0}".format(reason.decode("utf-8")))
                    return None, 0

                if proc.returncode == 1:
                    retry_reason = reason.decode("utf-8")

            except OSError as ex:
                if ex.errno != errno.ENOENT:
                    sys.stderr.write("\nCouldn't run test-failure-policy: {0}\n".format(str(ex)))

        # HACK: many tests are unstable, always retry them 3 times unless affected
        if not affected and not retry_reason:
            retry_reason = "be robust against unstable tests"

        has_unexpected_message = testlib.UNEXPECTED_MESSAGE.encode() in self.output
        has_pixel_test_message = testlib.PIXEL_TEST_MESSAGE.encode() in self.output
        if self.retries < 2 and not (has_unexpected_message or has_pixel_test_message) and retry_reason:
            self.retries += 1
            self._print_test(retry_reason="# RETRY {0} ({1})".format(self.retries, retry_reason))
            return retry_reason, 0

        self.output += b"\n"
        self._print_test()
        self.machine_id = None
        return None, 1

    # internal methods

    def __str__(self):
        return "{0} {1} {2}{3}".format(
            self.test_id,
            self.command[0],
            self.command[-1],
            " [ND@{0}]".format(self.machine_id) if self.nondestructive else "",
        )

    def _print_test(self, print_tap=True, retry_reason="", skip_reason=""):
        for line in self.output.strip().splitlines(keepends=True):
            while line:
                try:
                    sys.stdout.buffer.write(line)
                    break
                except BlockingIOError as e:
                    line = line[e.characters_written:]
                    time.sleep(0.1)

        if retry_reason:
            retry_reason = " " + retry_reason
        if skip_reason:
            skip_reason = " " + skip_reason

        if not print_tap:
            print(retry_reason + skip_reason)
            flush_stdout()
            return

        print()  # Tap needs to start on a separate line
        if self.returncode == 0:
            print(f"ok {self}{retry_reason}")
        elif skip_reason:
            print(f"ok {self}{skip_reason}")
        else:
            print(f"not ok {self}{retry_reason}")
        flush_stdout()


class GlobalMachine:
    def __init__(self, restrict=True, cpus=None, memory_mb=None):
        self.image = testvm.DEFAULT_IMAGE
        self.network = testvm.VirtNetwork(image=self.image)
        self.networking = self.network.host(restrict=restrict)
        self.machine = testvm.VirtMachine(verbose=True, networking=self.networking, image=self.image, cpus=cpus, memory_mb=memory_mb)
        if not os.path.exists(self.machine.image_file):
            self.machine.pull(self.machine.image_file)
        self.machine.start()
        self.start_time = time.time()
        self.duration = None
        self.ssh_address = f"{self.machine.ssh_address}:{self.machine.ssh_port}"
        self.web_address = f"{self.machine.web_address}:{self.machine.web_port}"
        self.running_test = None

    def reset(self):
        # It is important to re-use self.networking here, so that the
        # machine keeps its browser and control port.
        self.machine.kill()
        self.machine = testvm.VirtMachine(verbose=True, networking=self.networking, image=self.image)
        self.machine.start()

    def kill(self):
        assert self.running_test is None, "can't kill global machine with running test"
        self.machine.kill()
        self.network.kill()
        self.duration = round(time.time() - self.start_time)
        self.machine = None
        self.ssh_address = None
        self.web_address = None

    def is_available(self):
        return self.machine and self.running_test is None


def check_valid(filename):
    name = os.path.basename(filename)
    allowed = string.ascii_letters + string.digits + '-_'
    if not all(c in allowed for c in name):
        return None
    return name.replace("-", "_")


def build_command(filename, test, opts):
    cmd = [filename]
    if opts.trace:
        cmd.append("-t")
    if opts.verbosity:
        cmd.append("-v")
    if not opts.fetch:
        cmd.append("--nonet")
    if opts.list:
        cmd.append("-l")
    if opts.coverage:
        cmd.append("--coverage")
    cmd.append(test)
    return cmd


def get_affected_tests(test_dir, base_branch, test_files):
    if not base_branch:
        return []

    changed_tests = []

    # Detect affected tests from changed test files
    diff_out = subprocess.check_output(["git", "diff", "--name-only", "origin/" + base_branch, test_dir])
    # Never consider 'test/verify/check-example' to be affected - our tests for tests count on that
    # This file provides only examples, there is no place for it being flaky, no need to retry
    changed_tests = [test.decode("utf-8") for test in diff_out.strip().splitlines() if not test.endswith(b"check-example")]

    # If more than 3 test files were changed don't consider any of them as affected
    # as it might be a PR that changes more unrelated things.
    if len(changed_tests) > 3:
        # If 'test/verify/check-testlib' is affected, keep just that one - our tests for tests count on that
        if "test/verify/check-testlib" in changed_tests:
            changed_tests = ["test/verify/check-testlib"]
        else:
            changed_tests = []

    # Detect affected tests from changed pkg/* subdirectories in cockpit
    # If affected tests get detected from pkg/* changes, don't apply the
    # "only do this for max. 3 check-* changes" (even if the PR also changes ≥ 3 check-*)
    # (this does not apply to other projects)
    diff_out = subprocess.check_output(["git", "diff", "--name-only", "origin/" + base_branch, "--", "pkg/"])

    # Drop changes in css files - this does not affect tests thus no reason to retry
    files = [f.decode("utf-8") for f in diff_out.strip().splitlines() if not f.endswith(b"css")]

    changed_pkgs = set("check-" + pkg.split('/')[1] for pkg in files)
    changed_tests.extend([test for test in test_files if any(pkg in test for pkg in changed_pkgs)])

    return changed_tests


def detect_tests(test_files, image, opts):
    '''Build the list of tests we'll parallelize and the ones we'll run serially'''

    parallel_tests = []
    serial_tests = []
    seen_classes = {}
    test_id = 1

    for filename in test_files:
        name = check_valid(filename)
        if not name or not os.path.isfile(filename):
            continue
        loader = importlib.machinery.SourceFileLoader(name, filename)
        module = importlib.util.module_from_spec(importlib.util.spec_from_loader(loader.name, loader))
        loader.exec_module(module)
        for test_suite in unittest.TestLoader().loadTestsFromModule(module):
            for test in test_suite:
                # ensure that test classes are unique, so that they can be selected properly
                cls = test.__class__.__name__
                if seen_classes.get(cls) not in [None, filename]:
                    raise ValueError("test class %s in %s already defined in %s" % (cls, filename, seen_classes[cls]))
                seen_classes[cls] = filename

                test_method = getattr(test.__class__, test._testMethodName)
                test_str = "{0}.{1}".format(cls, test._testMethodName)
                # most tests should take much less than 10mins, so default to that;
                # longer tests can be annotated with @timeout(seconds)
                # check the test function first, fall back to the class'es timeout
                test_timeout = getattr(test_method, "__timeout", getattr(test, "__timeout", 600))
                if opts.tests and not any([t in test_str for t in opts.tests]):
                    continue
                if test_str in opts.exclude:
                    continue
                nd = getattr(test_method, "_testlib__non_destructive", False)
                rwa = getattr(test_method, "_testlib__retry_when_affected", True)
                test = Test(test_id, build_command(filename, test_str, opts), test_timeout, nd, rwa)
                if nd:
                    serial_tests.append(test)
                else:
                    if not opts.nondestructive:
                        parallel_tests.append(test)
                test_id += 1

    # sort serial tests by class/test name, to avoid spurious errors where failures depend on the order of
    # execution but let's make sure we always test them both ways around; hash the image name, which is
    # robust, reproducible, and provides an even distribution of both directions
    serial_tests.sort(key=lambda t: t.command[-1], reverse=bool(binascii.crc32(image.encode()) & 1))

    return (serial_tests, parallel_tests)


def list_tests(opts):
    test_files = glob.glob(os.path.join(opts.test_dir, "check-*"))
    serial_tests, parallel_tests = detect_tests(test_files, "dummy", opts)
    names = {t.command[-1] for t in serial_tests + parallel_tests}
    for n in sorted(names):
        print(n)


def run(opts, image):
    fail_count = 0
    start_time = time.time()

    if opts.coverage:
        # This gives us a convenient link at the top of the logs, see link-patterns.json
        print("Code coverage report in Coverage")

    test_files = glob.glob(os.path.join(opts.test_dir, "check-*"))
    changed_tests = get_affected_tests(opts.test_dir, opts.base, test_files)
    serial_tests, parallel_tests = detect_tests(test_files, image, opts)
    serial_tests_len = len(serial_tests)
    parallel_tests_len = len(parallel_tests)

    if opts.machine:
        assert not parallel_tests

    print(f"1..{serial_tests_len + parallel_tests_len}")
    flush_stdout()

    running_tests = []
    global_machines = []

    if not opts.machine:
        # Create appropriate number of serial machines; prioritize the nondestructive (serial) tests, to get
        # them out of the way as fast as possible, then let the destructive (parallel) ones start as soon as
        # a given serial runner is done.
        num_global = min(serial_tests_len, opts.jobs)

        for i in range(num_global):
            global_machines.append(GlobalMachine(restrict=not opts.enable_network, cpus=opts.nondestructive_cpus,
                                                 memory_mb=opts.nondestructive_memory_mb))

    # test scheduling loop
    while True:
        made_progress = False

        # mop up finished tests
        logging.debug(f"test loop: {len(running_tests)} running tests")
        for test in running_tests.copy():
            poll_result = test.poll()
            if poll_result is not None:
                made_progress = True
                running_tests.remove(test)
                test_machine = test.machine_id  # test_finish() resets it
                retry_reason, test_result = test.finish(changed_tests, print_tap=not opts.list, thorough=opts.thorough, track_naughties=opts.track_naughties)
                fail_count += test_result
                logging.debug(f"test {test} finished; result {test_result} retry reason {retry_reason}")

                if test_machine is not None and not opts.machine:
                    # unassign from global machine
                    global_machines[test_machine].running_test = None

                    # sometimes our global machine gets messed up; also, tests that time out don't run cleanup handlers
                    # restart it to avoid an unbounded number of test retries and follow-up errors
                    if not opts.machine and (poll_result == 124 or (retry_reason and "test harness" in retry_reason)):
                        # try hard to keep the test output consistent
                        sys.stderr.write("\nRestarting global machine %s\n" % test_machine)
                        sys.stderr.flush()
                        global_machines[test_machine].reset()

                # run again if needed
                if retry_reason:
                    if test.nondestructive:
                        serial_tests.insert(0, test)
                    else:
                        parallel_tests.insert(0, test)

        if opts.machine:
            if not running_tests and serial_tests:
                test = serial_tests.pop(0)
                logging.debug(f"Static machine is free, assigning next test {test}")
                test.assign_machine(-1, opts.machine, opts.browser)
                test.start()
                running_tests.append(test)
                made_progress = True
        else:
            # find free global machines, and either assign a new serial test, or kill them to free resources
            for (idx, machine) in enumerate(global_machines):
                if machine.is_available():
                    if serial_tests:
                        test = serial_tests.pop(0)
                        logging.debug(f"Global machine {idx} is free, assigning next test {test}")
                        machine.running_test = test
                        test.assign_machine(idx, machine.ssh_address, machine.web_address)
                        test.start()
                        running_tests.append(test)
                    else:
                        logging.debug(f"Global machine {idx} is free, and no more serial tests; killing")
                        machine.kill()

                    made_progress = True

        # fill the remaining available job slots with parallel tests
        while parallel_tests and len(running_tests) < opts.jobs:
            test = parallel_tests.pop(0)
            logging.debug(f"{len(running_tests)} running tests, starting next parallel test {test}")
            test.start()
            running_tests.append(test)
            made_progress = True

        # are we done?
        if not running_tests:
            assert not serial_tests, f"serial_tests should be empty: {[str(t) for t in serial_tests]}"
            assert not parallel_tests, f"parallel_tests should be empty: {[str(t) for t in parallel_tests]}"
            break

        # Sleep if we didn't make progress
        if not made_progress:
            time.sleep(0.5)

    # Create coverage report
    if opts.coverage:
        base_dir = os.path.realpath(f'{__file__}/../../..')
        output = os.environ.get("TEST_ATTACHMENTS", base_dir)
        lcov_files = glob.glob(f"{base_dir}/lcov/*.info.gz")
        if len(lcov_files) > 0:
            try:
                subprocess.check_call(["genhtml", "--no-function-coverage", "--output-dir", f"{output}/Coverage"] + lcov_files)
            except subprocess.CalledProcessError as e:
                print(f"Failed to create coverage report: {e}")

    # print summary
    duration = int(time.time() - start_time)
    hostname = socket.gethostname().split(".")[0]

    serial_details = []
    if not opts.machine:
        for (idx, machine) in enumerate(global_machines):
            serial_details.append(f"{idx}: {machine.duration}s")

    details = f"[{duration}s on {hostname}, {parallel_tests_len} parallel tests, {serial_tests_len} serial tests: {', '.join(serial_details)}]"
    print()
    if fail_count > 0:
        print(f"# {fail_count} TESTS FAILED {details}")
    else:
        print(f"# TESTS PASSED {details}")
    flush_stdout()

    return fail_count


def main():
    parser = testlib.arg_parser(enable_sit=False)
    parser.add_argument('-j', '--jobs', type=int,
                        default=int(os.environ.get("TEST_JOBS", 1)), help="Number of concurrent jobs")
    parser.add_argument('--thorough', action='store_true',
                        help='Thorough mode, no skipping known issues')
    parser.add_argument('-n', '--nondestructive', action='store_true',
                        help='Only consider @nondestructive tests')
    parser.add_argument('--machine', metavar="hostname[:port]",
                        default=None, help="Run tests against an already running machine;  implies --nondestructive")
    parser.add_argument('--browser', metavar="hostname[:port]",
                        default=None, help="When using --machine, use this cockpit web address")
    parser.add_argument('--test-dir', default=os.environ.get("TEST_DIR", testvm.TEST_DIR),
                        help="Directory in which to glob check-* files; default: %(default)s")
    parser.add_argument('--exclude', action="append", default=[], metavar="TestClass.testName",
                        help="Exclude test (exact match only); can be specified multiple times")
    parser.add_argument('--nondestructive-cpus', type=int, default=None,
                        help="Number of CPUs for nondestructive test global machines")
    parser.add_argument('--nondestructive-memory-mb', type=int, default=None,
                        help="RAM size for nondestructive test global machines")
    parser.add_argument('--base', default=os.environ.get("BASE_BRANCH"),
                        help="Retry affected tests compared to given base branch; default: %(default)s")
    parser.add_argument('--track-naughties', action='store_true',
                        help='Update the occurrence of naughties on cockpit-project/bots')
    opts = parser.parse_args()

    if opts.machine:
        if opts.jobs > 1:
            parser.error("--machine cannot be used with concurrent jobs")
        if not opts.browser:
            parser.error("--browser must be specified together with --machine")
        opts.nondestructive = True

    # Tell any subprocesses what we are testing
    if "TEST_REVISION" not in os.environ:
        r = subprocess.run(["git", "rev-parse", "HEAD"],
                           universal_newlines=True, check=False, stdout=subprocess.PIPE)
        if r.returncode == 0:
            os.environ["TEST_REVISION"] = r.stdout.strip()

    os.environ["TEST_BROWSER"] = os.environ.get("TEST_BROWSER", "chromium")

    image = testvm.DEFAULT_IMAGE
    testvm.DEFAULT_IMAGE = image
    os.environ["TEST_OS"] = image

    # Make sure tests can make relative imports
    sys.path.append(os.path.realpath(opts.test_dir))

    if opts.list:
        list_tests(opts)
        return 0

    return run(opts, image)


if __name__ == '__main__':
    # logging.basicConfig(level=logging.DEBUG)
    sys.exit(main())
